import requests
import json
import os
from constants import Constants



class OllamaAPI:
    def __init__(self, model_name=Constants.OLLAMA_API_MODEL):
        self.model_name = model_name
        self.BASE_URL = Constants.OLLAMA_API_URL

    def generate_text(self, prompt, options=None):
        url = f"{self.BASE_URL}/api/generate"
        data = {
            "model": self.model_name,
            "prompt": prompt,
            "stream": False  # 不使用流式输出
        }
        
        # 添加可选参数
        if options:
            data.update(options)
        
        try:
            # 发送POST请求
            response = requests.post(url, json=data)
            response.raise_for_status()  # 检查请求是否成功
            
            # 解析响应
            result = response.json()
            return result.get("response", "")
        
        except requests.exceptions.RequestException as e:
            print(f"请求出错: {e}")
            return None
        except json.JSONDecodeError as e:
            print(f"解析响应出错: {e}")
            return None
    
    def stop_model(self):
        """停止并释放模型占用的资源（如显存），但保留模型文件"""
        url = "/api/stop"
        data = {"name": self.model_name}
        response = requests.post(self.BASE_URL+url, json=data)
        
        if response.status_code == 200:
            print(f"模型 {self.model_name} 已停止，显存已释放")

    def __del__(self):
        # self.stop_model()
        pass
        


if __name__ == "__main__":
    t1 = "吴棠望着那女子美得惊心动魄的俏靥，不甘的眼神，摇摇欲坠的身影，纵使他早以在暗处见过多遍，但今天距离如此之近，也不免有几份惊艳。似乎察觉出那女子的不甘，吴棠不管这女子是否还有意识，轻叹道：“今天我从你这里拿走的，总有一天，我会加倍还给你的。"
    t2 = "写一首7律诗，以秋天、枫叶为主题"
    ollama_api = OllamaAPI(model_name="qwen3:8b")  # 替换为你想要使用的模型名称
    # prompt = """结合原文分析这句话的语速、音调、感情色彩。原文为：%s 。\n 分析的语句为：%s 。返回的格式为{语速：（1-5），语调：（1-5），感情色彩：（一个中文词）}"""%(t1,t2)
    
    options = {"temperature": 0.7, "top_p": 0.9}  # 可选参数
    # result = ollama_api.generate_text(t2, options)  # 调用API并获取生成的文本
    # print(result)




    with open(r"H:\ComfyUI_windows_portable\prompting-guide_wan1.3b.json","r",encoding="utf-8") as f :
        with open(r"H:\ComfyUI_windows_portable\prompting-guide_wan1.3b_cn.csv","w",encoding="utf-8") as f2:
            data = json.load(f)

            for item in data.items():
                # print(item)
                prompt = """将以下内容翻译为中文，直接返回结果：%s"""%(item[1])
                text = ollama_api.generate_text(prompt, options)  # 调用API并获取生成的文本
                result = text.split("</think>")[-1]
                print(result)

                f2.write(item[0]+","+result+"\n")